from pyspark.sql.session import SparkSession
from pyspark.ml.linalg import Vectors

from pyspark.ml.regression import LinearRegression

# 创建环境
spark = SparkSession.builder.getOrCreate()

# 1、读取数据
# label:标签，目标值，y
# features: 特征向量，x
# 向量，稠密向量，稀疏向量
data = spark.read.format("libsvm").load("../../data/poins.txt")

data.printSchema()
data.show()

# 2、选择算法
# LinearRegression: 线性回归算法
lr = LinearRegression()

# 3、将数据带入算法训练模型
# 内部需要通过不断循环迭代计算最优解
model = lr.fit(data)

print("k: %s" % str(model.coefficients))
print("b: %s" % str(model.intercept))

# 4、带入新的数据预测结果
# Vectors.dense(100.0): 构建向量
print(model.predict(Vectors.dense(100.0)))

# 5、保存模型
model.save("../../data/model")